import re
from WSOption import WSOption
from WSJob import WSJob

class MaskGenerator():
	name = 'MaskGenerator'
	mask = None

	options = {
		"mask": WSOption("mask", "Brute mask", "", True, ['--mask'])
	}

	#def __init__(self, mask):
	#	self.mask = mask

	def keys_in_str(self, str, keys):
		for key in keys:
			if str.find(key) != -1:
				return True
		return False

	def generate_masks(self, mask):
		charsets = {'?d': '0123456789',
					'?l': 'abcdefghijklmnopqrstuvwxyz',
					'?u': 'ABCDEFGHIJKLMNOPQRSTUVWXYZ',
					'?s': '\!@#$%^&*()[]{}`~-_=+|;:\'",.<>/?'}

		masks = {}
		if re.match('^.*,\d+,\d+$', mask):
			end_limit = int(mask[mask.rfind(',')+1:])
			tmp = mask[:mask.rfind(',')]
			start_limit = int(tmp[tmp.rfind(',')+1:])
			mask = tmp[:tmp.rfind(',')]
			for i in range(start_limit, end_limit+1):
				mask_name = mask + ',' + str(start_limit) + ',' + str(i)
				masks[mask_name] = {}
				for c in range(i):
					masks[mask_name][c+1] = self.raw_mask(mask, charsets)
		else:
			masks[mask] = {}
			mask_symbols = self.get_symbols(mask)
			for s in mask_symbols:
				masks[mask][s+1] = self.raw_mask(mask_symbols[s], charsets)

		return masks

	def raw_mask(self, mask, charsets):
		return_mask = []

		while self.keys_in_str(mask, charsets.keys()):
			for charset in charsets:
				mask = mask.replace(charset, charsets[charset])

		for s in mask:
			return_mask.append(s)

		return return_mask

	def get_symbols(self, mask):
		symbols = {}
		counter = 0
		while True:
			try:
				symbol = mask[0]
				if symbol == '?':
					symbols[counter] = symbol + mask[1]
					mask = mask[2:]
				else:
					symbols[counter] = symbol
					mask = mask[1:]
				counter += 1
			except IndexError:
				break

		return symbols

	def dict(self):
		masks = self.generate_masks(self.options['mask'].value)

		dicts = []
		for mask in masks:
			dict = []
			for symbolNum in masks[mask]:
				tmp = []
				for symbol in masks[mask][symbolNum]:
					if len(dict):
						for word in dict:
							tmp.append(word + symbol)
					else:
						tmp.append(symbol)
				dict = tmp
			dicts.extend(dict)
		return dicts

	def getJob(self):
		job = WSJob()
		for w in self.dict():
			job.put(w)
		return job